# PostgreSQL客户端模块 - 用于历史数据的按年分表存储

import psycopg2
import pandas as pd
import json
import logging
from typing import Dict, List, Any, Optional, Union, Tuple
from datetime import datetime, timedelta
from psycopg2.extras import execute_values
import numpy as np
import os
from src.utils.schemas import STOCK_PRICES_SCHEMA, TECHNICAL_INDICATORS_SCHEMA, STOCK_LIST_SCHEMA

logger = logging.getLogger(__name__)

class PostgresClient:
    """
    PostgreSQL客户端类，用于历史数据的按年分表存储
    实现按年分表的数据存储策略
    """
    
    def __init__(self, host: str = 'localhost', port: int = 5432,
                 database: str = 'smartnexus', user: str = 'postgres',
                 password: str = None):
        """
        初始化PostgreSQL连接
        
        Args:
            host: 数据库服务器地址
            port: 数据库服务器端口
            database: 数据库名称
            user: 用户名
            password: 密码
        """
        self._connected = False
        try:
            self.conn_params = {
                'host': host,
                'port': port,
                'database': database,
                'user': user,
                'password': password
            }
            
            # 连接到PostgreSQL
            self.conn = psycopg2.connect(**self.conn_params)
            self.conn.autocommit = False  # 使用事务管理
            self.test_connection()
            self._connected = True
            
            logger.info(f"PostgreSQL连接初始化成功: {host}:{port}/{database}")
            
            # 初始化表结构
            self._init_tables()
        except Exception as e:
            logger.error(f"PostgreSQL连接初始化失败: {str(e)}")
            raise

    def test_connection(self) -> bool:
        """测试PostgreSQL连接"""
        try:
            with self.conn.cursor() as cur:
                cur.execute('SELECT 1')
            self._connected = True
            return True
        except:
            self._connected = False
            return False
    
    def is_connected(self) -> bool:
        """获取连接状态"""
        if not self._connected:
            return self.test_connection()
        return self._connected
    
    def _init_tables(self):
        """初始化必要的表结构和分区表"""
        try:
            with self.conn.cursor() as cursor:
                # 检查表是否存在
                cursor.execute("""
                    SELECT tablename FROM pg_tables 
                    WHERE schemaname = 'public';
                """)
                existing_tables = {row[0] for row in cursor.fetchall()}

                # 定义表结构
                tables = {
                    'stock_prices': STOCK_PRICES_SCHEMA,
                    'technical_indicators': TECHNICAL_INDICATORS_SCHEMA,
                    'stock_list': STOCK_LIST_SCHEMA
                }

                # 动态创建表
                for table_name, schema in tables.items():
                    if (table_name not in existing_tables):
                        columns_def = [f"{col} {dtype}" for col, dtype in schema['columns'].items()]
                        primary_key = f", PRIMARY KEY ({', '.join(schema['primary_key'])})"
                        foreign_keys = "".join(
                            [f", FOREIGN KEY ({fk['column']}) REFERENCES {fk['references']} ON DELETE CASCADE" for fk in schema.get('foreign_keys', [])]
                        )
                            
                        # 添加唯一约束和分区
                        unique_constraints = ""
                        partition_clause = ""
                        if table_name == 'stock_prices':
                            unique_constraints = ", CONSTRAINT unique_symbol_start_date_period UNIQUE (id, start_date)"
                            partition_clause = "PARTITION BY RANGE (start_date)"
                        elif table_name == 'technical_indicators':
                            unique_constraints = ", CONSTRAINT unique_price_id_start_date UNIQUE (price_id, start_date)"
                            partition_clause = "PARTITION BY RANGE (start_date)"
                                
                        create_table_sql = f"""
                            CREATE TABLE {table_name} (
                                {', '.join(columns_def)}
                                {primary_key}
                                {foreign_keys}
                                {unique_constraints}
                            )
                            {partition_clause};
                        """
                        
                        cursor.execute(create_table_sql)
                        logger.info(f"创建表: {table_name}")

                        # 如果是 stock_prices 表，添加复合主键
                        # if table_name == 'stock_prices':
                        #     cursor.execute("""
                        #         ALTER TABLE stock_prices ADD PRIMARY KEY (id, start_date);
                        #     """)                        

                        # 创建索引
                        if 'indexes' in schema:
                            for index in schema['indexes']:
                                index_sql = f"""
                                    CREATE INDEX {index['name']} ON {table_name} ({', '.join(index['columns'])});
                                """
                                cursor.execute(index_sql)
                                logger.info(f"创建索引: {index['name']}")

                # 提交事务
                self.conn.commit()

                logger.info("PostgreSQL表结构初始化成功")
        except Exception as e:
            self.conn.rollback()
            logger.error(f"PostgreSQL表结构初始化失败: {str(e)}")
            raise

    def _get_partition_table_name(self, table_name: str, partition_type: str, date: datetime) -> str:
        """
        根据分区类型和日期生成分区表名。

        Args:
            table_name (str): 主表名。
            partition_type (str): 分区类型（'year' 或 'month'）。
            date (datetime): 数据日期。

        Returns:
            str: 分区表名。
        """
        if partition_type == 'year':
            return f"{table_name}_{date.year}"
        elif partition_type == 'month':
            return f"{table_name}_{date.year}_{date.month:02d}"
        else:
            raise ValueError(f"未知的分区类型: {partition_type}")

    def _ensure_partition_table(self, table_name: str, schema: Dict, date: datetime):
        """
        确保分区表存在。

        Args:
            table_name (str): 主表名。
            schema (Dict): 表的schema定义。
            date (datetime): 数据日期。
        """
        partition_type = schema.get('partition')
        if not partition_type:
            return  # 非分区表无需处理

        partition_table_name = self._get_partition_table_name(table_name, partition_type, date)

        try:
            with self.conn.cursor() as cursor:
                # 检查 PostgreSQL 版本是否支持分区表
                cursor.execute("SHOW server_version;")
                version = cursor.fetchone()[0]
                major_version = int(version.split('.')[0])
                if major_version < 10:
                    raise RuntimeError(f"PostgreSQL 版本 {version} 不支持分区表功能，请升级到 10 或更高版本")

                # 检查主表是否为分区表
                cursor.execute("""
                    SELECT relkind FROM pg_class WHERE relname = %s;
                """, (table_name,))
                result = cursor.fetchone()
                if result is None:
                    raise RuntimeError(f"主表 {table_name} 不存在")
                if result[0] != 'p':  # 'p' 表示分区表
                    # 将主表转换为分区表
                    logger.info(f"主表 {table_name} 不是分区表，正在转换为分区表")
                    cursor.execute(f"""
                        ALTER TABLE {table_name}
                        SET (
                            PARTITION BY RANGE (start_date)
                        );
                    """)

                # 检查分区表是否存在
                cursor.execute("""
                    SELECT 1 FROM pg_class c
                    JOIN pg_namespace n ON n.oid = c.relnamespace
                    WHERE c.relname = %s AND n.nspname = 'public';
                """, (partition_table_name,))

                if cursor.fetchone() is None:
                    # 创建分区表
                    if partition_type == 'year':
                        start_date = f"{date.year}-01-01"
                        end_date = f"{date.year + 1}-01-01"
                    elif partition_type == 'month':
                        start_date = f"{date.year}-{date.month:02d}-01"
                        next_month = date.replace(day=28) + timedelta(days=4)
                        end_date = next_month.replace(day=1).strftime('%Y%m%d')

                    cursor.execute(f"""
                        CREATE TABLE {partition_table_name} PARTITION OF {table_name}
                        FOR VALUES FROM ('{start_date}') TO ('{end_date}');
                    """)

                    logger.info(f"创建分区表: {partition_table_name}")

                self.conn.commit()
        except Exception as e:
            self.conn.rollback()
            logger.error(f"创建分区表失败: {str(e)}")
            raise

    def query_table(self, table_name: str, conditions: Optional[Dict[str, Any]] = None, 
                    columns: Optional[List[str]] = None, limit: Optional[int] = None, 
                    order_by: Optional[str] = None) -> pd.DataFrame:
        """
        通用的数据库查询方法。

        Args:
            table_name (str): 表名。
            conditions (Dict[str, Any], optional): 查询条件，键为列名，值为条件值。
            columns (List[str], optional): 要查询的列，默认为所有列。
            limit (int, optional): 限制返回的行数。
            order_by (str, optional): 排序条件，例如 'date DESC'。

        Returns:
            pd.DataFrame: 查询结果。
        """
        try:
            # 构建查询列
            select_columns = ', '.join(columns) if columns else '*'

            # 构建查询条件
            where_clause = ''
            params = []
            if conditions:
                where_conditions = [f"{col}  %s" for col in conditions.keys()]
                where_clause = f"WHERE {' AND '.join(where_conditions)}"
                params = list(conditions.values())

            # 构建排序和限制
            order_clause = f"ORDER BY {order_by}" if order_by else ''
            limit_clause = f"LIMIT {limit}" if limit else ''

            # 构建完整SQL
            query = f"""
                SELECT {select_columns} 
                FROM {table_name} 
                {where_clause} 
                {order_clause} 
                {limit_clause}
            """

            # 执行查询
            with self.conn.cursor() as cursor:
                cursor.execute(query, params)
                result = cursor.fetchall()
                columns = [desc[0] for desc in cursor.description]
                return pd.DataFrame(result, columns=columns)

        except Exception as e:
            logger.error(f"查询表 {table_name} 失败: {str(e)}")
            return pd.DataFrame()

    def store_data(self, table_name: str, data: pd.DataFrame) -> bool:
        """
        存储数据到指定表。

        Args:
            table_name (str): 表名。
            data (pd.DataFrame): 数据的 DataFrame 格式。

        Returns:
            bool: 存储是否成功。
        """
        try:
            if data.empty:
                logger.warning(f"表 {table_name} 的数据为空，跳过存储")
                return False

            # 获取表的 schema 定义
            schema = None
            if table_name == 'stock_prices':
                schema = STOCK_PRICES_SCHEMA
            elif table_name == 'technical_indicators':
                schema = TECHNICAL_INDICATORS_SCHEMA
            elif table_name == 'stock_list':
                schema = STOCK_LIST_SCHEMA

            # 如果是分区表，确保分区表存在
            if schema and 'partition' in schema:
                start_dates = data['start_date']
                if isinstance(start_dates, pd.Series):
                    dates = pd.to_datetime(start_dates).dt.to_pydatetime()
                else:
                    dates = [pd.to_datetime(start_dates).to_pydatetime()]
                for date in dates:
                    self._ensure_partition_table(table_name, schema, date)

            # 确保列名与数据库表一致
            columns = list(data.columns)
            values = data.to_records(index=False).tolist()

            # 动态生成 INSERT 语句
            insert_sql = f"""
                INSERT INTO {table_name} ({', '.join(columns)})
                VALUES %s
                ON CONFLICT DO NOTHING
            """

            # 执行批量插入
            with self.conn.cursor() as cursor:
                execute_values(cursor, insert_sql, values)
                self.conn.commit()

            logger.info(f"成功存储 {len(data)} 条数据到表 {table_name}")
            return True
        except Exception as e:
            self.conn.rollback()
            logger.error(f"存储数据到表 {table_name} 失败: {str(e)}")
            return False

    def update_data(self, table_name: str, data: Dict[str, Any], conditions: Dict[str, Any]) -> bool:
        """
        通用的更新数据方法。

        Args:
            table_name (str): 表名。
            data (Dict[str, Any]): 要更新的数据，键为列名，值为更新值。
            conditions (Dict[str, Any]): 更新条件，键为列名，值为条件值。

        Returns:
            bool: 更新是否成功。
        """
        try:
            if not data:
                logger.warning("更新数据为空，跳过更新操作")
                return False

            # 构建SET子句
            set_clause = ', '.join([f"{col} = %s" for col in data.keys()])
            set_values = list(data.values())

            # 构建WHERE子句
            where_clause = ' AND '.join([f"{col} %s" for col in conditions.keys()])
            where_values = list(conditions.values())

            # 构建完整SQL
            sql = f"""
                UPDATE {table_name}
                SET {set_clause}
                WHERE {where_clause}
            """

            # 执行SQL语句
            with self.conn.cursor() as cursor:
                cursor.execute(sql, set_values + where_values)
                self.conn.commit()

            logger.info(f"成功更新表{table_name}的数据")
            return True

        except Exception as e:
            self.conn.rollback()
            logger.error(f"更新表{table_name}的数据失败: {str(e)}")
            return False

    def delete_data(self, table_name: str, conditions: Dict[str, Any]) -> bool:
        """
        删除表中符合条件的数据。

        Args:
            table_name (str): 表名。
            conditions (Dict[str, Any]): 删除条件，键为列名，值为条件值。

        Returns:
            bool: 删除是否成功。
        """
        try:
            # 构建条件语句
            condition_str = ' AND '.join([f"{col} = %s" for col in conditions.keys()])
            values = list(conditions.values())

            sql = f"DELETE FROM {table_name} WHERE {condition_str}"

            with self.conn.cursor() as cursor:
                cursor.execute(sql, values)
                self.conn.commit()

            logger.info(f"成功删除表{table_name}中符合条件的数据")
            return True

        except Exception as e:
            self.conn.rollback()
            logger.error(f"删除表{table_name}数据失败: {str(e)}")
            return False

    def close(self):
        """
        关闭PostgreSQL连接
        """
        if hasattr(self, 'conn'):
            self.conn.close()
    def __del__(self):
        """
        析构函数，确保连接被关闭
        """
        self.close()